/**
 * This program is free software, you can redistribute it and/or modify it.
 * Copyright (c) 2025 Huawei Technologies Co., Ltd.
 * This file is a part of the CANN Open Software.
 * Licensed under CANN Open Software License Agreement Version 2.0 (the "License").
 * Please refer to the License for details. You may not use this file except in compliance with the License.
 * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, INCLUDING
 * BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE.
 * See LICENSE in the root of the software repository for the full text of the License.
 */

/*!
 * \file erf_tiling.cpp
 * \brief  AscendC custom erf op tiling implementation
 */
#include "register/op_impl_registry.h"
#include "tiling/platform/platform_ascendc.h"
#include "tiling_base/tiling_util.h"
#include "log/log.h"
#include "util/math_util.h"
#include "../op_kernel/erf_tiling_data.h"
#include "../op_kernel/erf_tiling_key.h"
#include "graph/utils/type_utils.h"

namespace optiling {
using namespace Ops::Math::OpTiling;

#define BLOCK_SIZE 512U
constexpr uint32_t BUFFER_NUM = 2;
constexpr uint32_t WS_SYS_SIZE = 0;
struct ErfCompileInfo {};

/* ---------- 1. 解析阶段（预留） ---------- */
static ge::graphStatus TilingParseForErf([[maybe_unused]] gert::TilingParseContext* context)
{
    OP_CHECK_IF(context == nullptr, OP_LOGE(context, "context is nullptr"), return ge::GRAPH_FAILED);
    return ge::GRAPH_SUCCESS;
}

/* ---------- 2. 取平台信息 ---------- */
static ge::graphStatus GetPlatformInfo(gert::TilingContext* context, uint64_t& ubSize, int64_t& coreNum)
{
    OP_CHECK_IF(context == nullptr, OP_LOGE(context, "context is nullptr"), return ge::GRAPH_FAILED);
    auto ascendcPlatform = platform_ascendc::PlatformAscendC(context->GetPlatformInfo());
    ascendcPlatform.GetCoreMemSize(platform_ascendc::CoreMemType::UB, ubSize);
    coreNum = ascendcPlatform.GetCoreNumAiv();
    // ASCENDXXX请替换为实际的版本型号

    OP_CHECK_IF(coreNum <= 0, OP_LOGE(context, "coreNum is 0"), return ge::GRAPH_FAILED);
    OP_CHECK_IF(ubSize <= 0, OP_LOGE(context, "ubSize is 0"), return ge::GRAPH_FAILED);
    return ge::GRAPH_SUCCESS;
}

static ge::graphStatus GetWorkspaceSize(gert::TilingContext* context)
{
    OP_CHECK_IF(context == nullptr, OP_LOGE(context, "context is nullptr"), return ge::GRAPH_FAILED);
    size_t usrSize = 0;
    auto ascendcPlatform = platform_ascendc::PlatformAscendC(context->GetPlatformInfo());
    uint32_t sysWorkspaceSize = ascendcPlatform.GetLibApiWorkSpaceSize();
    size_t* currentWorkspace = context->GetWorkspaceSizes(
        1); // 通过框架获取workspace的指针，GetWorkspaceSizes入参为所需workspace的块数。当前限制使用一块。
    currentWorkspace[0] = usrSize + sysWorkspaceSize;
    return ge::GRAPH_SUCCESS;
}

/* ---------- 4. 输入形状 & 数据类型 ---------- */
static ge::graphStatus GetShapeAttrsInfo([[maybe_unused]]ge::DataType& dataType,
    gert::TilingContext* context, uint64_t ubSize, uint64_t& inputNum, uint64_t& inputBytes,
    uint64_t& tileBlockNum,[[maybe_unused]]uint64_t tileBlockNum1,uint64_t& tileDataNum, uint64_t& inputLengthAlign512)
{
    OP_CHECK_IF(context == nullptr || context->GetInputShape(0) == nullptr,
                OP_LOGE(context, "invalid context"), return ge::GRAPH_FAILED);
    auto inputDesc = context->GetInputDesc(0);
    dataType = inputDesc->GetDataType();
    inputNum = context->GetInputShape(0)->GetStorageShape().GetShapeSize();
    uint32_t typeLen = 0;
    ge::TypeUtils::GetDataTypeLength(context->GetInputDesc(0)->GetDataType(), typeLen);
    uint64_t inputLen = inputNum * typeLen;
    if (inputNum == 0) return ge::GRAPH_FAILED;

    inputBytes = typeLen;
    inputLengthAlign512 = (((inputLen + BLOCK_SIZE - 1) / BLOCK_SIZE) * BLOCK_SIZE);
    double ubDataNumber1 = static_cast<uint64_t>(
    (typeLen == 4) ? 3.0 :
    (typeLen == 8) ? 3.5 : 8.0);
    double ubDataNumber1 = (typeLen == 4) ? 3.0 :
    (typeLen == 8) ? 3.5 : 8.0;

    double ubDataNumber = (typeLen == 4) ? 5.0 :
    (typeLen == 8) ? 6.5 : 10.0;
    /* 根据 dtype 决定 UB 中同时驻留 buffer 条数*/

    tileBlockNum = (ubSize / BLOCK_SIZE) / ubDataNumber;
    if (tileBlockNum == 0) return ge::GRAPH_FAILED;

    tileDataNum = (tileBlockNum * BLOCK_SIZE) / inputBytes;
    return ge::GRAPH_SUCCESS;
}



static ge::graphStatus CalculateCoreBlockNums(
    uint64_t inputLengthAlign512, int64_t coreNum,[[maybe_unused]]uint64_t tileBlockNum1, uint64_t tileBlockNum, uint64_t inputBytes,
    uint64_t tileDataNum, uint64_t& smallCoreDataNum, uint64_t& bigCoreDataNum,
    uint64_t& smallTailDataNum, uint64_t& bigTailDataNum, uint64_t& finalSmallTileNum,
    uint64_t& finalBigTileNum, uint64_t& tailBlockNum)
/* ---------- 5. 核间分块 ---------- */
{
    uint64_t inputNum = inputLengthAlign512 / inputBytes; 
    if (BLOCK_SIZE == 0 || coreNum == 0 || tileBlockNum == 0 || inputBytes == 0)
        return ge::GRAPH_FAILED;
    uint64_t tileDataNum1 = (tileBlockNum1 * BLOCK_SIZE) / inputBytes;
  /* 核数调整 */
    if (tileDataNum1 >= inputNum) {
        coreNum = 1;
    }
    uint64_t blkPerCore = inputLengthAlign512 / BLOCK_SIZE / coreNum;
    if(blkPerCore <= tileBlockNum1){
        tileBlockNum =tileBlockNum1;
        tileDataNum =tileDataNum1;}
    /* 核数调整 */
   
  
    tailBlockNum = (inputLengthAlign512 / BLOCK_SIZE) % coreNum;

    /* small core */
    smallCoreDataNum = blkPerCore * BLOCK_SIZE / inputBytes;
    uint64_t smallTile = blkPerCore / tileBlockNum;
    finalSmallTileNum = (blkPerCore % tileBlockNum == 0) ? smallTile : smallTile + 1;
    smallTailDataNum = smallCoreDataNum - tileDataNum * smallTile;
    smallTailDataNum = (smallTailDataNum == 0) ? tileDataNum : smallTailDataNum;

    /* big core */
    blkPerCore += 1;
    bigCoreDataNum = blkPerCore * BLOCK_SIZE / inputBytes;
    uint64_t bigTile = blkPerCore / tileBlockNum;
    finalBigTileNum = (blkPerCore % tileBlockNum == 0) ? bigTile : bigTile + 1;
    bigTailDataNum = bigCoreDataNum - tileDataNum * bigTile;
    bigTailDataNum = (bigTailDataNum == 0) ? tileDataNum : bigTailDataNum;

    return ge::GRAPH_SUCCESS;
}

/* ---------- 6. 主 tiling 函数 ---------- */
static ge::graphStatus ErfTilingFunc(gert::TilingContext* context)
{
    ErfTilingData* tiling = context->GetTilingData<ErfTilingData>();
    OP_CHECK_NULL_WITH_CONTEXT(context, tiling);
    OP_CHECK_IF(memset_s(tiling, sizeof(ErfTilingData), 0, sizeof(ErfTilingData)) != EOK,
                OP_LOGE(context, "set tiling data error"), return ge::GRAPH_FAILED);

    /* 平台信息 */
    uint64_t ubSize;
    int64_t  coreNum;
    auto ret = GetPlatformInfo(context, ubSize, coreNum);
    OP_CHECK_IF(ret != ge::GRAPH_SUCCESS, OP_LOGE(context, "GetPlatformInfo error"), return ge::GRAPH_FAILED);
    ge::DataType dataType;
    /* 输入信息 */
  
    uint64_t inputNum = 0, inputBytes = 0, tileBlockNum = 0, tileDataNum = 0, tileBlockNum1 = 0, inputLengthAlign512 = 0;
    ret = GetShapeAttrsInfo(dataType,context, ubSize, inputNum, inputBytes,
                            tileBlockNum,tileBlockNum1, tileDataNum, inputLengthAlign512);
    OP_CHECK_IF(ret != ge::GRAPH_SUCCESS, OP_LOGE(context, "GetShapeAttrsInfo error"), return ge::GRAPH_FAILED);


    /* 分块计算 */
    uint64_t smallCoreDataNum, bigCoreDataNum, smallTailDataNum, bigTailDataNum;
    uint64_t finalSmallTileNum, finalBigTileNum, tailBlockNum;
    ret = CalculateCoreBlockNums(inputLengthAlign512, coreNum, tileBlockNum1 ,tileBlockNum,
                                 inputBytes, tileDataNum, smallCoreDataNum, bigCoreDataNum,
                                 smallTailDataNum, bigTailDataNum, finalSmallTileNum,
                                 finalBigTileNum, tailBlockNum);
    OP_CHECK_IF(ret != ge::GRAPH_SUCCESS, OP_LOGE(context, "CalculateCoreBlockNums error"), return ge::GRAPH_FAILED);

    /* 写回 tiling 结构 */
    tiling->smallCoreDataNum  = static_cast<uint32_t>(smallCoreDataNum);
    tiling->bigCoreDataNum    = static_cast<uint32_t>(bigCoreDataNum);
    tiling->tileDataNum       = static_cast<uint16_t>(tileDataNum);
    tiling->smallTailDataNum  = static_cast<uint16_t>(smallTailDataNum);
    tiling->bigTailDataNum    = static_cast<uint16_t>(bigTailDataNum);
    tiling->finalSmallTileNum = static_cast<uint32_t>(finalSmallTileNum);
    tiling->finalBigTileNum   = static_cast<uint32_t>(finalBigTileNum);
    tiling->tailBlockNum      = static_cast<uint8_t>(tailBlockNum);

    /* workspace & key */
    OP_CHECK_IF(GetWorkspaceSize(context) != ge::GRAPH_SUCCESS,
                OP_LOGE(context, "GetWorkspaceSize error"), return ge::GRAPH_FAILED);
    /* ---------- 6.1 计算 tilingKey ---------- */
    uint64_t tilingKey = 0;

    if (dataType == ge::DT_FLOAT) {
        if (finalBigTileNum == 1) {
            tilingKey = GET_TPL_TILING_KEY(ELEMENTWISE_TPL_SCH_MODE_0);
        } else {
            tilingKey = GET_TPL_TILING_KEY(ELEMENTWISE_TPL_SCH_MODE_1);
        }
    } else if (dataType == ge::DT_FLOAT16 || dataType == ge::DT_BF16 || dataType == ge::DT_DOUBLE) {
        if (finalBigTileNum == 1) {
            tilingKey = GET_TPL_TILING_KEY(ELEMENTWISE_TPL_SCH_MODE_2);
        } else {
            tilingKey = GET_TPL_TILING_KEY(ELEMENTWISE_TPL_SCH_MODE_3);
        }
    } else {
        OP_LOGE(context, "Unsupported dtype[%d]", dataType);
        return ge::GRAPH_FAILED;
    }

    context->SetTilingKey(tilingKey);


    context->SetBlockDim(coreNum);
    return ge::GRAPH_SUCCESS;
}

/* ---------- 7. 注册 ---------- */
IMPL_OP_OPTILING(Erf)
    .Tiling(ErfTilingFunc)
    .TilingParse<ErfCompileInfo>(TilingParseForErf);

}  // namespace optiling